pub mod error;

use self::error::{CallError, SerOrDeserErr};
use crate::{
    courier::codec::*,
    crate_common_types::*,
    utils::{deser_from_decoded, maybe_deser_owned::VariantMaybeDeserOwned, OwnedOrRefWithSrc},
};
use fnv::FnvHashMap;
use futures::{ready, Sink, Stream};
use log::warn;
use pin_project::pin_project;
use serde::{Deserialize, Serialize};
use std::{
    future::Future,
    io::Error as IoErr,
    net::SocketAddr,
    pin::Pin,
    sync::{
        atomic::{AtomicU64, Ordering},
        Arc,
    },
    task::{Context, Poll},
    time::Duration,
};
use tokio::{
    io::{AsyncRead, AsyncWrite},
    net::TcpStream,
    sync::{mpsc, oneshot},
};
use tokio_util::codec::Framed;

macro_rules! proj {
    ($it: ident) => {
        $it.as_mut().project()
    };
}

pub type ClientTuple<Req, Resp, C = TcpStream, Ser = DefaultSer, Deser = DefaultDeser> = (
    ClientServing<Req, Resp, C, Ser, Deser>,
    RequestHandle<Req, Resp, Ser, Deser>,
);

pub async fn with_default_serder_tcp<Req, Resp>(
    addr: SocketAddr,
) -> Result<ClientTuple<Req, Resp>, IoErr> {
    Ok(with_default_serder(TcpStream::connect(addr).await?))
}

pub fn with_default_serder<C, Req, Resp>(c: C) -> ClientTuple<Req, Resp, C>
where
    C: AsyncWrite + AsyncRead,
{
    new_client(c, DefaultSer::default(), DefaultDeser::default())
}

pub fn new_client<C, Ser, Deser, Req, Resp>(
    c: C,
    ser: Ser,
    deser: Deser,
) -> ClientTuple<Req, Resp, C, Ser, Deser>
where
    C: AsyncRead + AsyncWrite,
    Ser: ToWriterSerializer,
    Deser: FromSliceDeserializer + Clone,
{
    let next_id = Arc::new(AtomicU64::new(0));
    let (call_chan_tx, call_chan_rx) = mpsc::channel(1000);
    let (cancel_chan_tx, cancel_chan_rx) = mpsc::unbounded_channel();
    let raw_end = RawChannelEnd::new(c, WithFixedSizedPrefix::new(ser, deser.clone()));
    let client_serving = ClientServing {
        raw_end,
        cancel_chan_rx,
        call_chan_rx,
        awaiting_reqs: Default::default(),
        cancel_chan_tx: cancel_chan_tx.clone(),
        call_chan_tx: call_chan_tx.clone(),
    };
    let req_handle = RequestHandle {
        next_id,
        call_chan_tx,
        cancel_chan_tx,
        deser,
    };
    (client_serving, req_handle)
}

const RESP_ONESHOT_RECV_LIMIT: usize = 100;
const REQ_ONESHOT_SEND_LIMIT: usize = 100;

type RawChannelEnd<C, S, D, I, O> =
    Framed<C, WithFixedSizedPrefix<S, D, ReqFramePrefix, RespFramePrefix, I, O>>;

type CallResp<R, SE, DE> = Result<(RespFramePrefix, Decoded<R>), SerOrDeserErr<SE, DE>>;

type RespChanTx<R, SE, DE> = oneshot::Sender<CallResp<R, SE, DE>>;

struct CallReqWithRespChan<Req, Resp, SE, DE> {
    id: ClientLocalReqId,
    req: Req,
    resp_chan_tx: RespChanTx<Resp, SE, DE>,
}

/// A RPC client manager connected to **one** RPC server.
#[pin_project]
pub struct ClientServing<Req, Resp, Conn, Ser: ToWriterSerializer, Deser: FromSliceDeserializer> {
    #[pin]
    raw_end: RawChannelEnd<Conn, Ser, Deser, Req, Resp>,
    cancel_chan_rx: mpsc::UnboundedReceiver<ClientLocalReqId>,
    call_chan_rx: mpsc::Receiver<CallReqWithRespChan<Req, Resp, Ser::Err, Deser::Err>>,
    awaiting_reqs: FnvHashMap<ClientLocalReqId, RespChanTx<Resp, Ser::Err, Deser::Err>>,

    cancel_chan_tx: mpsc::UnboundedSender<ClientLocalReqId>,
    call_chan_tx: mpsc::Sender<CallReqWithRespChan<Req, Resp, Ser::Err, Deser::Err>>,
}

enum PollReceiveDone {
    NothingToReceive,
    OneshotReceiveLimitReached,
}

enum PollSendDone {
    OneshotSendLimitReached,
}

impl<Req, Resp, Ser, Der, Conn> ClientServing<Req, Resp, Conn, Ser, Der>
where
    Req: Serialize + VariantMaybeDeserOwned,
    Resp: Deserialize<'static> + VariantMaybeDeserOwned,
    Ser: ToWriterSerializer,
    Der: FromSliceDeserializer,
    Conn: AsyncRead + AsyncWrite,
{
    // Receive:
    // 1. responses from the underlying data channel
    fn poll_recv(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<PollReceiveDone, IoErr>> {
        for _ in 0..RESP_ONESHOT_RECV_LIMIT {
            match ready!(proj!(self).raw_end.poll_next(cx)) {
                Some(Ok((pre, resp))) => {
                    let id = pre.id;
                    if let Some(chan) = proj!(self).awaiting_reqs.remove(&id) {
                        let _ = chan.send(Ok((pre, resp)));
                    }
                }
                Some(Err(e)) => match e {
                    IoOrDecodeErr::Io(e) => return Poll::Ready(Err(e)),
                    IoOrDecodeErr::PayloadDeser(pre, e) => {
                        if let Some(tx) = proj!(self).awaiting_reqs.remove(&pre.id) {
                            let _ = tx.send(Err(SerOrDeserErr::DeserErr(e)));
                        }
                    }
                    IoOrDecodeErr::PrefixRead(e) => {
                        // TODO: handle prefix read error
                        warn!("A frame from server has invalid prefix: {e:?}")
                    }
                    IoOrDecodeErr::ShouldDelegate => {
                        // TODO: handle should delegate error
                        warn!("Resp isn't DeserializeOwned, the sender should delegate deser")
                    }
                },
                None => return Poll::Ready(Ok(PollReceiveDone::NothingToReceive)),
            }
        }
        Poll::Ready(Ok(PollReceiveDone::OneshotReceiveLimitReached))
    }

    // Send:
    // 1. call requests
    // 2. cancel requests
    fn poll_send(
        mut self: Pin<&mut Self>,
        cx: &mut Context<'_>,
    ) -> Poll<Result<PollSendDone, IoErr>> {
        for _ in 0..REQ_ONESHOT_SEND_LIMIT {
            let to_send: Option<(ReqFramePrefix, MaybeEmpty<Req>)> = 'to_send: {
                ready!(proj!(self).raw_end.poll_ready(cx)).map_err(IoOrEncodeErr::assert_io_err)?;
                loop {
                    match proj!(self).cancel_chan_rx.poll_recv(cx) {
                        Poll::Ready(None) => {
                            unreachable!("Unexpected state: closed cancel req channel!")
                        }
                        Poll::Ready(Some(id)) => {
                            if proj!(self).awaiting_reqs.remove(&id).is_none() {
                                continue;
                            }
                            break 'to_send Some((
                                ReqFramePrefix::new_cancel(id),
                                MaybeEmpty::Empty,
                            ));
                        }
                        Poll::Pending => break,
                    }
                }
                loop {
                    match proj!(self).call_chan_rx.poll_recv(cx) {
                        Poll::Ready(None) => {
                            unreachable!("Unexpected state: closed call req channel!")
                        }
                        Poll::Ready(Some(CallReqWithRespChan {
                            id,
                            req,
                            resp_chan_tx,
                        })) => {
                            if resp_chan_tx.is_closed() {
                                continue;
                            }
                            proj!(self).awaiting_reqs.insert(id, resp_chan_tx);
                            break 'to_send Some((
                                ReqFramePrefix::new_call(id, !req.variant_is_deser_owned()),
                                MaybeEmpty::Filled(req),
                            ));
                        }
                        Poll::Pending => break None,
                    }
                }
            };
            let Some(to_send) = to_send else {
                ready!(proj!(self).raw_end.poll_flush(cx)).map_err(IoOrEncodeErr::assert_io_err)?;
                return Poll::Pending;
            };
            let id = to_send.0.id;
            match proj!(self).raw_end.start_send(to_send) {
                Err(IoOrEncodeErr::Io(e)) => return Poll::Ready(Err(e)),
                Err(IoOrEncodeErr::PrefixWrite(_)) => unreachable!("Unexpected prefix error"),
                Err(IoOrEncodeErr::PayloadSer(e)) => {
                    let _ = proj!(self)
                        .awaiting_reqs
                        .remove(&id)
                        .unwrap()
                        .send(Err(SerOrDeserErr::SerErr(e)));
                }
                Ok(_) => {}
            }
        }
        Poll::Ready(Ok(PollSendDone::OneshotSendLimitReached))
    }
}

pub enum ClientServingDone {
    OtherEndClosedWithPendingReqs(usize),
}

impl<Req, Resp, Ser, Der, Conn> Future for ClientServing<Req, Resp, Conn, Ser, Der>
where
    Req: Serialize + VariantMaybeDeserOwned,
    Resp: Deserialize<'static> + VariantMaybeDeserOwned,
    Ser: ToWriterSerializer,
    Der: FromSliceDeserializer,
    Conn: AsyncRead + AsyncWrite,
{
    type Output = Result<ClientServingDone, IoErr>;

    fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Self::Output> {
        loop {
            match (self.as_mut().poll_send(cx)?, self.as_mut().poll_recv(cx)?) {
                (_, Poll::Ready(PollReceiveDone::NothingToReceive)) => {
                    break Poll::Ready(Ok(ClientServingDone::OtherEndClosedWithPendingReqs(
                        self.awaiting_reqs.len(),
                    )));
                }
                (Poll::Pending, Poll::Pending) => break Poll::Pending,
                _ => {}
            }
        }
    }
}

pub struct RequestHandle<Req, Resp, Ser: ToWriterSerializer, Der: FromSliceDeserializer> {
    next_id: Arc<AtomicU64>,
    call_chan_tx: mpsc::Sender<CallReqWithRespChan<Req, Resp, Ser::Err, Der::Err>>,
    cancel_chan_tx: mpsc::UnboundedSender<ClientLocalReqId>,
    deser: Der,
}

impl<Req, Resp, Ser, Der> Clone for RequestHandle<Req, Resp, Ser, Der>
where
    Ser: ToWriterSerializer,
    Der: FromSliceDeserializer,
    Der: Clone,
{
    fn clone(&self) -> Self {
        Self {
            next_id: Arc::clone(&self.next_id),
            call_chan_tx: self.call_chan_tx.clone(),
            cancel_chan_tx: self.cancel_chan_tx.clone(),
            deser: self.deser.clone(),
        }
    }
}

impl<Req, Resp, Ser, Der> RequestHandle<Req, Resp, Ser, Der>
where
    Ser: ToWriterSerializer,
    Der: FromSliceDeserializer,
    Req: Serialize + VariantMaybeDeserOwned,
    Resp: Deserialize<'static> + VariantMaybeDeserOwned,
{
    pub async fn call(
        &self,
        req: Req,
        timeout: Option<Duration>,
    ) -> Result<OwnedOrRefWithSrc<Resp>, CallError<Ser::Err, Der::Err>> {
        struct CancelGuard<'a, T> {
            do_cancel: bool,
            id: ClientLocalReqId,
            cancel_chan_tx: &'a mpsc::UnboundedSender<ClientLocalReqId>,
            resp_chan_rx: &'a mut oneshot::Receiver<T>,
        }
        impl<'a, T> Drop for CancelGuard<'a, T> {
            fn drop(&mut self) {
                if self.do_cancel {
                    self.resp_chan_rx.close();
                    let _ = self.cancel_chan_tx.send(self.id);
                }
            }
        }

        let (tx, mut rx) = oneshot::channel::<CallResp<Resp, Ser::Err, Der::Err>>();
        let id = self.next_id.fetch_add(1, Ordering::Relaxed) as ClientLocalReqId;

        self.call_chan_tx
            .send(CallReqWithRespChan {
                id,
                req,
                resp_chan_tx: tx,
            })
            .await
            .map_err(|_| CallError::LocalReqRecverDropped)?;

        let mut guard = CancelGuard {
            id,
            do_cancel: true,
            cancel_chan_tx: &self.cancel_chan_tx,
            resp_chan_rx: &mut rx,
        };

        let resp = if let Some(time) = timeout {
            match tokio::time::timeout(time, &mut guard.resp_chan_rx).await {
                Err(_) => return Err(CallError::TimeOut),
                Ok(resp) => resp,
            }
        } else {
            (&mut guard.resp_chan_rx).await
        };

        guard.do_cancel = false;

        let (pre, resp) = match resp {
            Err(_) => return Err(CallError::LocalRespSenderDropped),
            Ok(Err(e)) => return Err(e.into()),
            Ok(Ok(it)) => it,
        };
        match pre.typ {
            RespType::Done => {
                deser_from_decoded(resp, &self.deser).map_err(CallError::RespDeserErr)
            }
            RespType::Panicked => Err(CallError::CallPanic),
        }
    }
}
